/**
 * @file dog_control_for_Go1.cpp
 * @author 小沐
*/

#include "unitree_legged_sdk/unitree_legged_sdk.h"

// c++ 
#include <math.h>
#include <iostream>
#include <unistd.h>
#include <string.h>
#include <fstream>
#include <mutex>
#include <signal.h>

// ROS
#include <ros/ros.h>
#include <geometry_msgs/Twist.h>
#include <nav_msgs/Odometry.h>
#include <ros/package.h>
#include <std_msgs/Time.h>
#include <std_msgs/UInt8.h>
#include <std_msgs/Bool.h>
#include <sensor_msgs/Imu.h>

// 自定义
#include <dog_control/dog_info.h>

// tf
#include <tf/transform_broadcaster.h>

using namespace UNITREE_LEGGED_SDK;

// #define DEBUG                                        // 编译开关   开关debug 接口

///////////////////////////////////////////////////////////////////////////////
// 配置设置
const int AX_CONTROL_LEVEL = HIGHLEVEL;              // 控制模式
const LeggedType AX_LEG_TYPE = LeggedType::Go1;      // 设置控制机器人的类型
const uint16_t AX_LOCAL_PORT = 8090;                 // 设置本地发送端口
const char *AX_TARGET_IP = "192.168.123.161";        // 目标ip
const uint16_t AX_TARGET_PORT = 8082;                // 目标ip的端口
///////////////////////////////////////////////////////////////////////////////


///////////////////////////////////////////////////////////////////////////////
// 定义 速度发送消息格式
typedef struct
{
  double stamp;
  uint8_t vel_cmd_mode; 
  double vx;
  double vy;
  double vyaw;
} StampedVelocityCommand;
///////////////////////////////////////////////////////////////////////////////

///////////////////////////////////////////////////////////////////////////////
// 欧拉角
struct EulerAngles {
    double roll, pitch, yaw;
};

// 四元数
struct Quaternion
{
    double w, x, y, z;
};
///////////////////////////////////////////////////////////////////////////////


// 主程序

/**
 * Custom   机器人控制主程序，控制类
*/
class Custom
{
public:

  Safety safe_;


  UDP udp_;

  HighCmd cmd_ = {0};              

  HighState state_ = {0};

  // StampedVelocityCommand vel_cmd_;     // 用于转成SDK格式

  dog_control::dog_info vel_cmd_;

  // ROS

  std_msgs::Header  send_header;       // ROS的时间戳格式

  ros::NodeHandle nh_;                 // 句柄

  // 设置一些 topic 的发布器和订阅器

  ros::Subscriber sub_cmd_vel_;        // 订阅者 订阅cmd_vel，用于处理成SDK数据发出

  ros::Publisher dog_imu_pub;          // 发布者，发布imu数据

  ros::Publisher dog_odom_pub;         // 发布者，发布odom数据

  // 创建狗 imu 的数据

  sensor_msgs::Imu dogImu;             // 设置imu数据

  nav_msgs::Odometry dogOdom;          // 设置odom数据

  bool IsPubTf = false;                // 是否发布坐标系

  float dt_ = 400;

  tf::TransformBroadcaster odom_broadcaster;  // 创建一个tf 发布odom的tf

public:
  Custom(ros::NodeHandle nh,
          int control_level,
          LeggedType leg_type,
          uint16_t local_port,
          const char *target_ip,
          uint16_t target_port)
          :safe_(leg_type),udp_(local_port, target_ip, target_port, sizeof(HighCmd), sizeof(HighState))  // 冒号之后是赋值
  {
    // 原始控制相关
    velcmd_init();
    udp_.InitCmdData(cmd_);

    // ROS
    nh_ = nh;

    ROS_INFO("\033[32m pub_tf = %d",IsPubTf);

    ROS_INFO("\033[32m pub_frequency = %.2f",dt_);

    // 订阅速度指令
    sub_cmd_vel_ = nh.subscribe<geometry_msgs::Twist>("cmd_vel", 2, &Custom::cmdVelHandler, this); // for navigation
    dog_imu_pub = nh_.advertise<sensor_msgs::Imu>("dog_imu_raw", 1000);
    dog_odom_pub = nh_.advertise<nav_msgs::Odometry>("dog_odom", 1000);

  }

  void velcmd_init()
  {
    vel_cmd_.stamp = 0;
    vel_cmd_.mode = 0;
    vel_cmd_.gait_type = 0;
    vel_cmd_.speed_level = 0;
    vel_cmd_.foot_raise_height = 0;
    vel_cmd_.position[0] = 0;
    vel_cmd_.position[1] = 0;
    vel_cmd_.euler[0] = 0;
    vel_cmd_.euler[1] = 0;
    vel_cmd_.euler[2] = 0;
    vel_cmd_.velocity[0] = 0;
    vel_cmd_.velocity[1] = 0;
    vel_cmd_.yaw_speed = 0;
  }

  void walk(const double Vx , const double Vy ,const double Yaw){
    velcmd_init();
    vel_cmd_.stamp = ros::Time::now().toSec();
    vel_cmd_.mode = 2;    // mode = 2 行走 状态
    vel_cmd_.gait_type = 1;   //  gait_type = 1 正常 gait_type = 3 爬楼
    vel_cmd_.velocity[0] = Vx;
    vel_cmd_.velocity[1] = Vy;
    vel_cmd_.yaw_speed = Yaw;
  }

  void cmdVelHandler(const geometry_msgs::Twist::ConstPtr &msg)
  {
    walk(msg->linear.x,msg->linear.y,msg->angular.z);
  }

  void UDPRecv()
  {
    udp_.Recv();          // save data to buffer
    udp_.GetRecv(state_); // convert data type
    send_header.stamp = ros::Time::now();//更新发送的header ,用于赋值给imu和里程topic 的时间，可以保证之间的一致
    sendDogOdom();
    sendDogIMU();
  }

  void UDPSend()
  {
    updateCommand();
    udp_.SetSend(cmd_);
    udp_.Send();
  }

  EulerAngles ToEulerAngles(Quaternion q) 
  {
    EulerAngles angles;

    // roll (x-axis rotation)
    double sinr_cosp = 2 * (q.w * q.x + q.y * q.z);
    double cosr_cosp = 1 - 2 * (q.x * q.x + q.y * q.y);
    angles.roll = std::atan2(sinr_cosp, cosr_cosp);

    // pitch (y-axis rotation)
    double sinp = 2 * (q.w * q.y - q.z * q.x);
    if (std::abs(sinp) >= 1)
        angles.pitch = std::copysign(M_PI / 2, sinp); // use 90 degrees if out of range
    else
        angles.pitch = std::asin(sinp);

    // yaw (z-axis rotation)
    double siny_cosp = 2 * (q.w * q.z + q.x * q.y);
    double cosy_cosp = 1 - 2 * (q.y * q.y + q.z * q.z);
    angles.yaw = std::atan2(siny_cosp, cosy_cosp);

    return angles;
  }

  void sendDogOdom()
  {
    //四足机器人本体的里程记
    dogOdom.header.frame_id = "odom";
    dogOdom.child_frame_id = "base_link";
    dogOdom.header.stamp = send_header.stamp;

    // odom 位置信息
    dogOdom.pose.pose.position.x = state_.position[0];
    dogOdom.pose.pose.position.y = state_.position[1];
    dogOdom.pose.pose.position.z = state_.position[2]; 

    // odom 四元数使用imu的四元数
    dogOdom.pose.pose.orientation.w = state_.imu.quaternion[0];
    dogOdom.pose.pose.orientation.x = state_.imu.quaternion[1];
    dogOdom.pose.pose.orientation.y = state_.imu.quaternion[2];
    dogOdom.pose.pose.orientation.z = state_.imu.quaternion[3];
    
    //机器人在各个方向的速度反馈
    dogOdom.twist.twist.linear.x = state_.velocity[0]; // 前进速度
    dogOdom.twist.twist.linear.y = state_.velocity[1]; // 侧面平移速度
    dogOdom.twist.twist.linear.z = state_.velocity[2]; // 上升下降速度

    // 角速度使用imu的数据
    dogOdom.twist.twist.angular.x = state_.imu.gyroscope[0];
    dogOdom.twist.twist.angular.y = state_.imu.gyroscope[1];
    dogOdom.twist.twist.angular.z = state_.imu.gyroscope[2];
    
    // 发布tf 
    if (IsPubTf)
    {
      try{
        geometry_msgs::TransformStamped odom_trans;
        
        odom_trans.header.frame_id = "odom";
        odom_trans.child_frame_id = "base_link";
        odom_trans.header.stamp =send_header.stamp;
        odom_trans.transform.translation.x = dogOdom.pose.pose.position.x;
        odom_trans.transform.translation.y = dogOdom.pose.pose.position.y;
        odom_trans.transform.translation.z = dogOdom.pose.pose.position.z;
        odom_trans.transform.rotation.w = dogOdom.pose.pose.orientation.w;
        odom_trans.transform.rotation.x = dogOdom.pose.pose.orientation.x;
        odom_trans.transform.rotation.y = dogOdom.pose.pose.orientation.y;
        odom_trans.transform.rotation.z = dogOdom.pose.pose.orientation.z;

        odom_broadcaster.sendTransform(odom_trans);
      }
      catch (tf::TransformException ex){
        ROS_ERROR("%s",ex.what());
        ros::Duration(1.0).sleep();
      }
    }

    // 发布里程计
    dog_odom_pub.publish(dogOdom);

    Quaternion q;
    q.w = dogOdom.pose.pose.orientation.w;
    q.x = dogOdom.pose.pose.orientation.x;
    q.y = dogOdom.pose.pose.orientation.y;
    q.z = dogOdom.pose.pose.orientation.z;
    EulerAngles e = ToEulerAngles(q);

    // 编译 调试 
    #ifdef DEBUG
      std::cout<< " mode = " << (int)state_.mode <<std::endl;
      std::cout<< " gaitType = " << (int)state_.gaitType <<std::endl;
      // std::cout<< " roll = " << e.roll <<std::endl;
      // std::cout<< " pitch = " << e.pitch <<std::endl;
      // std::cout<< " yaw = " << e.yaw <<std::endl;
    #endif
    
  }


  void sendDogIMU()
  {
    // 构建机器狗 imu 的数据
    dogImu.header.frame_id = "dog_imu_link";
    dogImu.header.stamp = send_header.stamp;

    // dog's orientation imu四元数信息  归一化的四元数
    dogImu.orientation.w = state_.imu.quaternion[0];
    dogImu.orientation.x = state_.imu.quaternion[1];
    dogImu.orientation.y = state_.imu.quaternion[2];
    dogImu.orientation.z = state_.imu.quaternion[3];
    
    // dog's angular_velocity 陀螺仪，角速度，原始数据
    dogImu.angular_velocity.x = state_.imu.gyroscope[0];
    dogImu.angular_velocity.y = state_.imu.gyroscope[1];
    dogImu.angular_velocity.z = state_.imu.gyroscope[2];

    // dog's linear_acceleration 加速度计，加速度，原始数据
    dogImu.linear_acceleration.x = state_.imu.accelerometer[0];
    dogImu.linear_acceleration.y = state_.imu.accelerometer[1];
    dogImu.linear_acceleration.z = state_.imu.accelerometer[2];

    // 协方差
    dogImu.linear_acceleration_covariance[0]=state_.velocity[0];
    dogImu.linear_acceleration_covariance[1]=state_.velocity[1];
    dogImu.linear_acceleration_covariance[2]=state_.velocity[2];

    dog_imu_pub.publish(dogImu);  
  }

  /**
   * @return true if there is valid command that need to send.
  */
  void updateCommand()
  {
    // 重置全部控制指令

    // 当前时间戳
    ros::Time time_cur = ros::Time::now();
    double stamp_cur = time_cur.toSec();
  
    // 融合控制指令
    if (time_cur.toSec() - vel_cmd_.stamp < 0.2)
    {
      cmd_.mode = vel_cmd_.mode;
      cmd_.gaitType = vel_cmd_.gait_type;
      // cmd_.bodyHeight = 0.07; 
      cmd_.velocity[0] = vel_cmd_.velocity[0];
      cmd_.velocity[1] = vel_cmd_.velocity[1];
      cmd_.yawSpeed = vel_cmd_.yaw_speed;
      std::cout << "vx = " << cmd_.velocity[0] << ", vy = " << cmd_.velocity[1] << ", vyaw = " << cmd_.yawSpeed << std::endl;
    }
    else{
      velcmd_init();
    }
  }
};

void MySigintHandler(int sig)
{
  ROS_INFO("shutting down!");
  ros::shutdown();
  exit(0);
}

int main(int argc, char **argv)
{

  ros::init(argc, argv, "dog_control_node");
  ros::NodeHandle nh("~");

  // 重载SIGINT Handler，监听Ctrl+C，来强制结束此程序
  signal(SIGINT, MySigintHandler);

  ROS_INFO("[Unitree Go1 CTR] Node dog control has started.");
  std::cout << "[Unitree Go1 Control] Control level is set to HIGH-level." << std::endl;
  std::cout << "[Unitree Go1 ] Make sure the robot is standing on the ground." << std::endl;

  Custom custom(nh, AX_CONTROL_LEVEL, AX_LEG_TYPE,
                AX_LOCAL_PORT, AX_TARGET_IP, AX_TARGET_PORT);

  InitEnvironment();

  LoopFunc loop_udpSend("udp_send",1/custom.dt_, 3, boost::bind(&Custom::UDPSend, &custom));
  LoopFunc loop_udpRecv("udp_recv",1/custom.dt_, 3, boost::bind(&Custom::UDPRecv, &custom));

  loop_udpSend.start();
  loop_udpRecv.start();

  ros::Rate loop_rate(custom.dt_);
  
  while (nh.ok())
  {
    ros::spinOnce();
    loop_rate.sleep();
  }

  return 0;
}
